#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdlib.h>
#include <limits.h>
#include <float.h>

#include "pub_lib.h"
#include "ray_trace.h"

struct Intersection
{
	bool hit;
	float distance;
	uint32_t index;
};

struct Bounds3 Union(const struct Bounds3* b1, const struct Bounds3* b2)
{
    struct Bounds3 ret;
	ret.pMin = get_min(b1->pMin, b2->pMin);
	ret.pMax = get_max(b1->pMax, b2->pMax);
    return ret;
}

//if a <= b, return 1; otherwise return 0
int obj_compare(int a, int b, enum splitAxis axis)
{
	struct Triangle* pa = gTriangle + a;
	struct Triangle* pb = gTriangle + b;
	if(axis == AXIS_X) {
		if(pa->center.x <= pb->center.x)
			return 1;
		return 0;
	}

	if(axis == AXIS_Y) {
		if(pa->center.y <= pb->center.y)
			return 1;
		return 0;
	}

	if(axis == AXIS_Z) {
		if(pa->center.z <= pb->center.z)
			return 1;
		return 0;
	}
	return 0;
}

// Standard partition process of QuickSort().
// It considers the last element as pivot
// and moves all smaller element to left of
// it and greater elements to right
int partition(int* arr, int l, int r, enum splitAxis axis)
{
	int x = arr[r], i = l;
	for (int j = l; j <= r - 1; j++) {
		//if (arr[j] <= x) {
		if(obj_compare(arr[j], x, axis)) {
			swapi(arr[i], arr[j]);
			i++;
		}
	}
	swapi(arr[i], arr[r]);
	return i;
}

// This function returns k'th smallest
// element in arr[l..r] using QuickSort
// based method. ASSUMPTION: ALL ELEMENTS
// IN ARR[] ARE DISTINCT
// return index of arr, val = arr[index]
int kthSmallest(int* arr, int l, int r, int k, enum splitAxis axis)
{
	// If k is smaller than number of
	// elements in array
	if (k > 0 && k <= r - l + 1) {

		// Partition the array around last
		// element and get position of pivot
		// element in sorted array
		int index = partition(arr, l, r, axis);

		// If position is same as k
		if (index - l == k - 1)
			return index;

		// If position is more, recur
		// for left subarray
		if (index - l > k - 1)
			return kthSmallest(arr, l, index - 1, k, axis);

		// Else recur for right subarray
		return kthSmallest(arr, index + 1, r, k - index + l - 1, axis);
	}

	// If k is more than number of elements in array
	return 0; //error
}

enum splitAxis get_longest(struct Bounds3* bound)
{
	Vector3 d = Vector3Subtract(bound->pMax, bound->pMin);
	if (d.x > d.y && d.x > d.z)
		return AXIS_X;
	if (d.y > d.z)
		return AXIS_Y;
	return AXIS_Z;
}

void splite_two(struct BVHBuildNode* node, int objCnt)
{
	enum splitAxis axis = node->splitAxis;

	int index = kthSmallest(node->objAry, 0, objCnt - 1, objCnt/2, axis);
	int middle = node->objAry[index];
	//split 2 subs: left <= middle, and right > middle
	int arysize = objCnt/2 + 2;
	node->leftAry = malloc(sizeof(int) * arysize);
	node->rightAry = malloc(sizeof(int) * arysize);
	node->leftCnt = 0;
	node->rightCnt = 0;
	for(int i = 0; i < objCnt; i++){
		if (obj_compare(node->objAry[i], middle, axis)) {
			node->leftAry[node->leftCnt] = node->objAry[i];
			node->leftCnt++;
		} else {
			node->rightAry[node->rightCnt] = node->objAry[i];
			node->rightCnt++;
		}//end if-else
	}//end for
}

//build BVH tree
struct BVHBuildNode* recursiveBuild(struct MeshTriangle *p, int* root_ary, int objs)
{
	struct BVHBuildNode* node = calloc(1, sizeof(struct BVHBuildNode));
	list_add_tail(&node->list, &p->bvh_nodes); //add to list for free

	int objCnt = objs;
	node->nPrimitives = objCnt;

	//node is root, first call
	if (!root_ary) {//init root array
		node->objAry = malloc(sizeof(int) * objCnt);
		for(int i = 0; i < objCnt; i++)
			node->objAry[i] = i;
	}else {
		node->objAry = root_ary;
	}

	//calc bound box
	struct Bounds3 bounds;
	struct Triangle* pa;
	pa = gTriangle + node->objAry[0];
	bounds = pa->box;
	for (int i = 1; i < objCnt; i++){
		pa = gTriangle + node->objAry[i];
		bounds = Union(&bounds, &pa->box);
	}
	node->bounds = bounds;
	if (objCnt <= p->bvh.minPrims) {//reach to leaf
		return node;
	}

	//find longest dimension
	node->splitAxis = get_longest(&bounds);
	splite_two(node, objCnt);
	//recrusive build left and right child
    node->left = recursiveBuild(p, node->leftAry, node->leftCnt);
    node->right = recursiveBuild(p, node->rightAry, node->rightCnt);

	return node;
}

void buildBVH(struct MeshTriangle *p)
{
	INIT_LIST_HEAD(&p->bvh_nodes);
	p->bvh.splitMethod = NAIVE;
	p->bvh.minPrims = 5;

	if (!p->numTriangles)
		return;

	struct BVHAccel *bvh = &p->bvh;
	bvh->root = recursiveBuild(p, NULL, p->numTriangles);
}

//free array and node
void destroyBVH(struct MeshTriangle *p)
{
	struct BVHBuildNode *cur;
	struct BVHBuildNode* n;
	//list_for_each_entry_safe(pos, n, head, member, type)
	list_for_each_entry_safe(cur, n, &p->bvh_nodes, list, struct BVHBuildNode) {
		list_del(&cur->list); //remove this node from list
		if(cur->objAry)
			free(cur->objAry);
		free(cur);
	}
}

//t = (point - orig) / dir
//update Tenter and Texit
void calc_time(float p_min, float p_max, float orig, float d, float* enter, float* exit)
{
	float t_min, t_max;

	//(point - orig) * invDir; invDir = 1/dir
	t_min = (p_min - orig) * d;
	t_max = (p_max - orig) * d;
	if (t_min > t_max)
		swapf(t_min, t_max);
	//enter = max(t_min); exit = min(t_max)
	*enter = max(t_min, *enter);
	*exit =  min(t_max, *exit);
}

//box intersect with ray, cal tmin and tmax
//return true if intersect, otherwise return false
bool box_intersection(Vector3 orig, Vector3 invDir, struct Bounds3* box)
{
	float tEnter = -FLT_MAX;
	float tExit = FLT_MAX;

	calc_time(box->pMin.x, box->pMax.x, orig.x, invDir.x, &tEnter, &tExit);
	calc_time(box->pMin.y, box->pMax.y, orig.y, invDir.y, &tEnter, &tExit);
	calc_time(box->pMin.z, box->pMax.z, orig.z, invDir.z, &tEnter, &tExit);
	return (tEnter < tExit) && (tExit > 0);
}

//triangle intersect with ray
struct Intersection triangle_i(int *ary, int cnt, const Vector3 orig, const Vector3 dir)
{
	bool intersect = false;
	float t, u, v;
	struct Intersection ret;
	float tnear = FLT_MAX;
	uint32_t index = 0;

	struct Triangle *p;
	for (int k = 0; k < cnt; k++)
	{
		p = gTriangle + ary[k];
		if (rayTriangleIntersect(p->v0, p->v1, p->v2,
			orig, dir, &t, &u, &v) && t < tnear){
			tnear = t;
			index = ary[k]; //ary[k] is real index!
			intersect |= true;
		}
	}
	ret.hit = intersect;
	ret.distance = tnear;
	ret.index = index;

	return ret;
}

/*
 * a hit, b hit:
 *		distance of b < a, chooseA = 0
 *		otherwise chooseA = 1;
 * a hit, b not hit, chooseA = 1
 * a not hit, b hit, chooseA = 0
 * a not hit, b not hit, chooseA = 1(a, b is same, both not hit)
 * so, only two cases: chooseA = 0
*/
int chooseA(struct Intersection* a, struct Intersection* b)
{
	int chooseA = 1;
	if(a->hit && b->hit) {
		if(b->distance <= a->distance)
			chooseA = 0;
	}else if(b->hit) {
		chooseA = 0;
	}
	return chooseA;
}

//Traverse the BVH to find intersection
//if intersect, return ture, and fill payload
//otherwise return false
struct Intersection bvh_intersection(struct BVHBuildNode* node, Vector3 orig,
	Vector3 dir, Vector3 invDir)
{
	struct Intersection res;
	res.hit = false;
	// if not intersect with bounding box->there cannot be real intersection
	if (!box_intersection(orig, invDir, &node->bounds))
		return res; //not hit

	// if this node is leaf node, get intersection with actually mode--not bounding box
	if (node->left == NULL && node->right == NULL)
	{
		//test intersection with all objects, and return closest intersection
		return triangle_i(node->objAry, node->nPrimitives, orig, dir);
	}

	struct Intersection a, b;
	a = bvh_intersection(node->left,  orig, dir, invDir);
	b = bvh_intersection(node->right, orig, dir, invDir);
	if(chooseA(&a, &b))
		return a;
	return b;
}

bool bvh_intersect(struct Object* obj,
	const Vector3 orig, const Vector3 dir,
	float* tnear, uint32_t* index, Vector2* uv)
{
	struct MeshTriangle *p = (struct MeshTriangle*)obj;
	if (!p->bvh.root)
		return false;

	Vector3 invDir;
	invDir.x = 1.0f / dir.x;
	invDir.y = 1.0f / dir.y;
	invDir.z = 1.0f / dir.z;

	struct Intersection res;
	res = bvh_intersection(p->bvh.root, orig, dir, invDir);
	*tnear = res.distance;
	*index = res.index;
	return res.hit;
}
